#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
@Time    : 24-5-24 下午15:37
@Author  : zhangweihao
@File    : tasks.py
'''

import datetime
import logging
from flask import g, session, request, current_app
from flask_restful import Resource
from model.auto_model.abnormal_check import AbnormalPrivacyAiModelModel,AbnormalPrivacyDataModel
from ext.responses import RespHelper
from .form import *
# from .serializas import *
from ext.aberrant import Error
import json
from sqlalchemy import func, and_
import pandas as pd
from io import BytesIO
from flask import make_response
import os
from configs.config import Configs
from uuid import uuid4
from celery_task.AI.tasks import start_train, get_summarize, get_analysis, get_top5
from celery import group
from .serializas import AbnormalPrivacyAiModelSchema,AbnormalPrivacyDataModelSchema,AbnormalPrivacyDataAnalysisSchema

class AIModelView(Resource):

    # 模型训练
    def post(self):
        form = AImodelForm().validatas_for_api()
        model = g.db.session.query(AbnormalPrivacyAiModelModel).filter(AbnormalPrivacyAiModelModel.model_type == form.train_type.data).first()
        if model.is_training == 1:
            return RespHelper.error(f'训练中，请勿重复训练')
        start_train.delay(form.train_type.data)
        with g.db.auto_commit():
            model.is_training = 1
            g.db.session.flush()
        return RespHelper.success(f'{form.train_type.data}开始训练')
    
    def get(self):
        data = {}
        model = g.db.session.query(AbnormalPrivacyAiModelModel).order_by(AbnormalPrivacyAiModelModel.order.asc()).all()
        if model:
            data = AbnormalPrivacyAiModelSchema(many=True).dump(model)
        return RespHelper.success(message='查询模型成功',data=data)

class AbnormalPrivacyDataView(Resource):
    def get(self):
        form = AbnormalPrivacyDataGetForm().validatas_for_api()
        filters = []
        if form.destination_ip.data:
            filters.append(AbnormalPrivacyDataModel.destination_ip.like(f'%{form.destination_ip.data}%'))
        if form.type.data:
            types = json.loads(form.type.data)
            filters.append(AbnormalPrivacyDataModel.type.in_(types))
        if form.handle_status.data:
            filters.append(AbnormalPrivacyDataModel.handle_status == int(form.handle_status.data))
        start_time = form.start_time.data
        end_time = form.end_time.data
        if start_time and end_time:
            filters.append(AbnormalPrivacyDataModel.time.between(start_time, end_time))
        source_ips = form.source_ips.data
        if source_ips:
            source_ip_list = json.loads(source_ips)
            ip_filter_conditions = []
            for ip in source_ip_list:
                ip_filter_conditions.append(AbnormalPrivacyDataModel.source_ips.contains(ip))
            filters.append(and_(*ip_filter_conditions))

        data = {}
        paginat = g.db.session.query(AbnormalPrivacyDataModel).filter(*filters).order_by(AbnormalPrivacyDataModel.time.desc()).paginate()
        if paginat:
            data = AbnormalPrivacyDataModelSchema(many=True).dump(paginat.items)
        return RespHelper.success(message='查询异常数据成功',data={'data':data, 'pagination': paginat.page_info})
        
    # 修改处理状态
    def post(self):
        form = AbnormalPrivacyDataPostForm().validatas_for_api()
        set_status = form.set_status.data
        ids = request.json.get('ids', [])
        handle_des = form.remark.data
        origin_data = g.db.session.query(AbnormalPrivacyDataModel).filter(AbnormalPrivacyDataModel.id.in_(ids)).all()
        destination_ips = []
        if origin_data:
            with g.db.auto_commit():
                for i in origin_data:
                    i.handle_status = set_status
                    if handle_des:
                        i.handle_des = handle_des
                    destination_ips.append(f'{i.destination_ip}')
                g.db.session.flush()
        return RespHelper.success(message=f'处理接收ip：【{",".join(destination_ips)}】状态成功')

class AbnormalPrivacyDataAnalysisView(Resource):
    def get(self):
        form = AbnormalPrivacyDataGetForm().validatas_for_api()
        obj = {
            'destination_ip':form.destination_ip.data,
            'type':form.type.data,
            'handle_status':form.handle_status.data,
            'start_time':form.start_time.data,
            'end_time':form.end_time.data,
            'source_ips':form.source_ips.data,
            'page':form.page.data,
            'limit':form.limit.data,
        }
        analysis_type = request.args.get('analysis','destination_ip')
        top5_type = request.args.get('top5','destination_ip')
        task_group = group(get_summarize.s(obj, TRAIN_TYPE_LIST), get_analysis.s(obj, analysis_type), get_top5.s(obj, top5_type))()
        task_group.join()

        return RespHelper.success(data={'summarize': task_group.results[0].get(), 'analysis':task_group.results[1].get(),
                                         'top5': task_group.results[2].get()}, sorted=False)

        
        
